package pers.tz.题目.leetcode.树;

import pers.tz.common.TreeNode;

import java.util.*;

/**
 * @author irony
 * @date 2024-03-14
 */
public class leetcode_0863_二叉树中所有距离为K的结点 {
//给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 k 。
//
// 返回到目标结点 target 距离为 k 的所有结点的值的列表。 答案可以以 任何顺序 返回。
//
// 示例 1：
//输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
//输出：[7,4,1]
//解释：所求结点为与目标结点（值为 5）距离为 2 的结点，值分别为 7，4，以及 1
//
// 示例 2:
//输入: root = [1], target = 1, k = 3
//输出: []
//
// 提示:
// 节点数在 [1, 500] 范围内
// 0 <= Node.val <= 500
// Node.val 中所有值 不同
// 目标结点 target 是树上的结点。
// 0 <= k <= 1000
    public static void main(String[] args) {
        TreeNode t1 = new TreeNode(3);
        TreeNode t2 = new TreeNode(5);
        TreeNode t3 = new TreeNode(1);
        TreeNode t4 = new TreeNode(6);
        TreeNode t5 = new TreeNode(2);
        TreeNode t6 = new TreeNode(0);
        TreeNode t7 = new TreeNode(8);
        TreeNode t8 = new TreeNode(7);
        TreeNode t9 = new TreeNode(4);

        t1.left = t2;
        t1.right = t3;

        t2.left = t4;
        t2.right = t5;

        t3.left = t6;
        t3.right = t7;

        t5.left = t8;
        t5.right = t9;

        System.out.println(distanceK(t1, t2, 2));
    }

    // 1、通过哈希表维护每个节点的父节点，将树发散为图
    // 2、通过广度优先搜索，找到路径=K的节点
    public static List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        // 子 - 父
        Map<TreeNode, TreeNode> parentMap = new HashMap<>();
        parentMap.put(root, null);
        generateParentMap(root, parentMap);

        List<Integer> ans = new ArrayList<>();
        Set<TreeNode> visited = new HashSet<>();
        visited.add(target);

        // 广度优先遍历
        Queue<TreeNode> queue = new LinkedList<>();
        queue.add(target); // 以target为起点发散
        int curLevel = 0;
        while (!queue.isEmpty()) {
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                TreeNode node = queue.poll();
                if (curLevel == k) {
                    ans.add(node.val);
                }

                if (node.left != null && !visited.contains(node.left)) {
                    visited.add(node.left);
                    queue.add(node.left);
                }
                if (node.right != null && !visited.contains(node.right)) {
                    visited.add(node.right);
                    queue.add(node.right);
                }
                // 头
                if (parentMap.get(node) != null && !visited.contains(parentMap.get(node))) {
                    visited.add(parentMap.get(node));
                    queue.add(parentMap.get(node));
                }
            }
            curLevel++;
            if (curLevel > k) {
                return ans;
            }
        }

        return ans;
    }

    private static void generateParentMap(TreeNode root, Map<TreeNode, TreeNode> parentMap) {
        if (root == null) {
            return;
        }

        if (root.left != null) {
            parentMap.put(root.left, root);
            generateParentMap(root.left, parentMap);
        }
        if (root.right != null) {
            parentMap.put(root.right, root);
            generateParentMap(root.right, parentMap);
        }
    }
}
